/**
 * @file hs_cuda_elemwise_mimo.cuh
 * @author iam002
 * @brief 并行应用场景: 逐元素运算 elementwise（多入多出 MIMO）
 * @version 1.0
 * @date 2024-10-15
 * 
 * @copyright Copyright (c) 2024
 * 
 */
#ifndef __HS_CUDA_ELEMWISE_MIMO_CUH__
#define __HS_CUDA_ELEMWISE_MIMO_CUH__

#include "hs_cuda_utils.cuh"

#ifdef HS_USE_CXX_STANDARD_11
#include "hs_c11ext.h"
#endif

namespace hs
{

namespace cuda
{

namespace elemwise
{

namespace mimo
{


/**************************************************************
调用链:
    Op<X>in<X>out (顶层封装)
        |
        |
    GenericLauncher
        |
        |___ launchKernel
                |
                |___ applyGeneric (__global__)
                        |
                        |___ applyPack/Tail
                                |
                                |___ Functor (由用户设计功能单元)

功能单元:
    1. 在 device 端重载()运算符, 前部分参数为输入变量, 后部分为输出变量, 输出变量使用&, 函数返回值为 void;
       例如: __device__ void operator(int a, int b, int c, float & out1, float & out2) const;
    2. 功能单元中的每个类成员必须有明确确定大小, 不要定义功能单元的类成员为:
        (a) std::vector 等标准库模板类
        (b) 指向 host 端地址的指针或指针数组
       个人猜测: 当功能单元传入核函数时, 会执行 cudaMalloc 和 cudaMemcpy 的操作, 在 device 端上
       复制一个同样的功能单元, 由于是 cudaMemcpy 是按照指定大小进行值的复制, 因此不建议成员中定义
       指向主机端地址的指针哈!


参考功能算子（输出变量使用引用）：

template<typename T>
struct ResolveFunctor {
    __device__ void operator()(T x, T y, T & z1, T & z2) const
    {
        z1 = (x + y) / 2.0;
        z2 = (x - y) / 2.0;
    }
};

***************************************************************/

/* 解包, 执行功能单元, 展开输出指针数组和可变长输入参数 */
template<size_t pack_size, typename FunctorType, typename OutputType, typename... InputType, std::size_t... Is>
__device__ void applyPack(const FunctorType & functor
    , Packed<pack_size, OutputType>** p_out_pack_list
    , const Packed<pack_size, InputType>*... p_in_pack
#ifdef HS_USE_CXX_STANDARD_11
    , hs::c11ext::index_sequence<Is...>
#else
    , std::index_sequence<Is...>
#endif
    )
{
    #pragma unroll
    for (int i = 0; i < pack_size; ++i) {
        functor((p_in_pack->elem[i])..., (p_out_pack_list[Is]->elem[i])...);
    }
}

/* 处理尾部数据, 同样执行功能单元, 展开输出指针数组和可变长输入参数  */
template<typename FunctorType, typename OutputType, typename... InputType, std::size_t... Is>
__device__ void applyTail(const FunctorType & functor
    , OutputType** p_out_list
    , const InputType*... p_in
#ifdef HS_USE_CXX_STANDARD_11
    , hs::c11ext::index_sequence<Is...>
#else
    , std::index_sequence<Is...>
#endif
    )
{
    functor((*p_in)..., (*(p_out_list[Is]))... );
}

/* 将指针数组中的每个指针前进 step 步 */
#ifdef HS_USE_CXX_STANDARD_11
// 辅助递归函数，用来逐个对列表元素进行操作
template<std::size_t I, typename T>
__device__ void forwardOneElement(T* src_list, T* dst_list, int step)
{
    dst_list[I] = src_list[I] + step;
}

// 基础函数模板，用于终止递归
template<typename T>
__device__ void forwardSomeStepImpl(T* src_list, T* dst_list, int step) {}

// 递归函数模板，用于展开参数包
template<typename T, std::size_t First, std::size_t... Rest>
__device__ void forwardSomeStepImpl(T* src_list, T* dst_list, int step) {
    forwardOneElement<First>(src_list, dst_list, step);  // 对当前索引的元素进行操作
    forwardSomeStepImpl<T, Rest...>(src_list, dst_list, step);  // 对剩余的索引递归展开
}

// 替代折叠表达式的主函数，使用 index_sequence 展开
template<typename T, std::size_t... Is>
__device__ void forwardSomeStep(T* src_list, T* dst_list, int step, hs::c11ext::index_sequence<Is...>) {
    forwardSomeStepImpl<T, Is...>(src_list, dst_list, step);  // 展开索引并调用递归展开函数
}
#else 
template<typename T, std::size_t... Is>
__device__ void forwardSomeStep(T* src_list, T* dst_list, int step
#ifdef HS_USE_CXX_STANDARD_11
    , hs::c11ext::index_sequence<Is...>
#else
    , std::index_sequence<Is...>
#endif
)
{
    ((dst_list[Is] = src_list[Is] + step), ...);
}
#endif

/* 核函数, 请注意, 核函数形参涉及到数组和指针的传递时, 应该保证它们是在device端, 这真的很重要;
 * 另外, 核函数也不要使用引用.
 */
template<size_t pack_size, int num_output, typename FunctorType, typename OutputType, typename... InputType>
__global__ void __launch_bounds__(KERNEL_BLOCK_SIZE)
applyGeneric(FunctorType functor
    , int num_pack
    , Packed<pack_size, OutputType>* p_out_pack_list[num_output]
    , const Packed<pack_size, InputType>*... p_in_pack
    , int num_tail
    , OutputType* p_out_tail_list[num_output]
    , const InputType*... p_in_tail)
{
    const int global_tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (int i = global_tid; i < num_pack; i += blockDim.x * gridDim.x) {
        /* 数组中的每个指针前进 i 步, 指向对应的输出数据地址 */
        Packed<pack_size, OutputType>* p_clone_out_pack_list[num_output];
        forwardSomeStep(p_out_pack_list
            , p_clone_out_pack_list
            , i
        #ifdef HS_USE_CXX_STANDARD_11
            , hs::c11ext::make_index_sequence<num_output>{}
        #else
            , std::make_index_sequence<num_output>{}
        #endif
        );
        
        /* 处理包 */
        applyPack<pack_size, FunctorType, OutputType, InputType...>(functor
            , p_clone_out_pack_list
            , (p_in_pack + i)...
        #ifdef HS_USE_CXX_STANDARD_11
            , hs::c11ext::make_index_sequence<num_output>{}
        #else
            , std::make_index_sequence<num_output>{}
        #endif
            );
    }
    if (global_tid < num_tail) {
        /* 同样, 移动输出数组中的每个指针 */
        forwardSomeStep(p_out_tail_list
            , p_out_tail_list
            , global_tid
        #ifdef HS_USE_CXX_STANDARD_11
            , hs::c11ext::make_index_sequence<num_output>{}
        #else
            , std::make_index_sequence<num_output>{}
        #endif
            );

        /* 处理尾部数据 */
        applyTail<FunctorType, OutputType, InputType...>(functor
            , p_out_tail_list
            , (p_in_tail + global_tid)...
        #ifdef HS_USE_CXX_STANDARD_11
            , hs::c11ext::make_index_sequence<num_output>{}
        #else
            , std::make_index_sequence<num_output>{}
        #endif
            );
    }
}

/* 准备工作, 用于启动核函数 */
template<int pack_size, int num_output, typename FunctorType, typename OutputType, typename... InputType>
cudaError_t launchKernel(const FunctorType & fucntor, size_t N, OutputType* p_out_list[num_output], const InputType*... p_in)
{
    /* 计算打包数和尾部剩余数据个数 */
    const int num_pack = N / pack_size;
    const int tail_offset = num_pack * pack_size;
    const int num_tail = N - tail_offset;

    /* 计算线程块数 */
    int num_blocks;
    {
        cudaError_t err = GetNumBlocks(num_pack, &num_blocks);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* NOTE 指针数组都在 host 端 */
    OutputType* p_out_tail_list[num_output];
    for (int i = 0; i < num_output; ++i) {
        p_out_tail_list[i] = p_out_list[i] + tail_offset;
    }

    Packed<pack_size, OutputType>* p_out_pack_list[num_output];
    for (int i = 0; i < num_output; ++i) {
        p_out_pack_list[i] = reinterpret_cast<Packed<pack_size, OutputType>*>(p_out_list[i]); 
    }

    /* FIXME 当向核函数传递指针数组时, 由于数组名就是地址, 实际传递参数只复制了这个地址,
     * 若这个地址在 host 端, 则在 device 端的核函数时不能获取真实的值, 因此必须在 device
     * 端复制一个同样的指针数组.
     */

    /* NOTE 在 device 端拷贝指针数组 */
    OutputType** p_device_out_tail_list = nullptr;
    {
        cudaError_t err = cudaMalloc(&p_device_out_tail_list, sizeof(p_out_tail_list));
        if (err != cudaSuccess) {
            return err;
        }
        err = cudaMemcpy(p_device_out_tail_list, p_out_tail_list, sizeof(p_out_tail_list), cudaMemcpyHostToDevice);
        if (err != cudaSuccess) {
            return err;
        }
    }

    Packed<pack_size, OutputType>** p_device_out_pack_list = nullptr;
    {
        cudaError_t err = cudaMalloc(&p_device_out_pack_list, sizeof(p_out_pack_list));
        if (err != cudaSuccess) {
            return err;
        }
        err = cudaMemcpy(p_device_out_pack_list, p_out_pack_list, sizeof(p_out_pack_list), cudaMemcpyHostToDevice);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 调用核函数 */
    applyGeneric<pack_size, num_output, FunctorType, OutputType, InputType...>
    <<<num_blocks, KERNEL_BLOCK_SIZE, 0>>>
    (
        fucntor
        , num_pack
        , p_device_out_pack_list
        , (reinterpret_cast<const Packed<pack_size, InputType>*>(p_in))...
        , num_tail
        , p_device_out_tail_list
        , (p_in + tail_offset) ...
    );

    /* 记得释放 device 端分配的内存 */
    cudaFree(p_device_out_tail_list);
    cudaFree(p_device_out_pack_list);

    return cudaPeekAtLastError();
}

/* 再次封装, 隐藏 pack_size */
template<int num_output, typename FunctorType, typename OutputType, typename... InputType>
struct GenericLauncher
{
    static cudaError_t launch(const FunctorType & functor, size_t N, OutputType* p_out_list[num_output], const InputType*... p_in)
    {
        constexpr int pack_size = PackSize<OutputType, InputType...>();
        
        cudaError_t err;

        OutputType* tmp = nullptr;
        if (isAlignedForPack<pack_size, OutputType, InputType...>(tmp, p_in...)) {
            err = launchKernel<pack_size, num_output, FunctorType, OutputType, InputType...>(functor, N, p_out_list, p_in...);
        }
        else {
            err = launchKernel<1, num_output, FunctorType, OutputType, InputType...>(functor, N, p_out_list, p_in...);
        }
        return err;
    }
};

/**
 * @brief 运算模板: 2输入,2输出
 * 
 * @tparam FunctorType 自定义功能单元结构体
 * @tparam OutputType 输出数据类型(所有输出类型必须都相同)
 * @tparam InputTypeA 输入数据类型(输入之间类型可不同)
 * @tparam InputTypeB 输入数据类型(输入之间类型可不同)
 * @param N 数组长度
 * @param p_in_a 在 device 端输入数据指针
 * @param p_in_b 在 device 端输入数据指针
 * @param functor 自定义功能单元结构体
 * @param p_out_1 在 device 端输出数据指针(保存结果)
 * @param p_out_2 在 device 端输出数据指针(保存结果)
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB> 
inline cudaError_t Op2in2out(size_t N
    , const InputTypeA* p_in_a
    , const InputTypeB* p_in_b
    , const FunctorType & functor
    , OutputType* p_out_1
    , OutputType* p_out_2)
{
    constexpr int num_output = 2;
    OutputType* p_out_list[num_output] = {p_out_1, p_out_2};
    return GenericLauncher<num_output, FunctorType, OutputType, InputTypeA, InputTypeB>::launch(functor, N, p_out_list, p_in_a, p_in_b);
}


/**
 * @brief 运算模板: 3输入,2输出
 * 
 * @tparam FunctorType 自定义功能单元结构体
 * @tparam OutputType 输出数据类型(所有输出类型必须都相同)
 * @tparam InputTypeA 输入数据类型(输入之间类型可不同)
 * @tparam InputTypeB 输入数据类型(输入之间类型可不同)
 * @tparam InputTypeC 输入数据类型(输入之间类型可不同)
 * @param N 数组长度
 * @param p_in_a 在 device 端输入数据指针
 * @param p_in_b 在 device 端输入数据指针
 * @param functor 自定义功能单元结构体
 * @param p_out_1 在 device 端输出数据指针(保存结果)
 * @param p_out_2 在 device 端输出数据指针(保存结果)
 * @return cudaError_t 
 */
template<typename FunctorType, typename OutputType, typename InputTypeA, typename InputTypeB, typename InputTypeC> 
inline cudaError_t Op3in2out(size_t N
    , const InputTypeA* p_in_a
    , const InputTypeB* p_in_b
    , const InputTypeC* p_in_c
    , const FunctorType & functor
    , OutputType* p_out_1
    , OutputType* p_out_2)
{
    constexpr int num_output = 2;
    OutputType* p_out_list[num_output] = {p_out_1, p_out_2};
    return GenericLauncher<num_output, FunctorType, OutputType, InputTypeA, InputTypeB, InputTypeC>::launch(functor, N, p_out_list, p_in_a, p_in_b, p_in_c);
}

} // end namespace mimo

} // end namespace elementwise

} // end namespace cuda

} // end namespace hs

#endif

